# 指定文件编码为UTF-8
# coding: utf-8

"""
optimizer_compare_naive.py 功能说明：
1. 可视化比较不同优化算法在简单二维函数上的优化轨迹
2. 测试的优化器包括：SGD、Momentum、AdaGrad、Adam
3. 使用函数 f(x,y) = x²/20 + y² 作为优化目标
4. 绘制各优化器的参数更新路径和函数等高线图
5. 直观展示不同优化器的收敛特性

关键点：
- 使用简单的二次函数清晰展示优化过程
- 四种优化器使用相同的初始点(-7.0, 2.0)
- 每种优化器迭代30次
- 通过参数更新路径观察优化行为
- 等高线图展示函数形状
"""

# 导入系统模块
import sys, os
# 添加父目录到系统路径
sys.path.append(os.pardir)
# 导入NumPy数值计算库
import numpy as np
# 导入matplotlib绘图库
import matplotlib.pyplot as plt
# 导入有序字典类
from collections import OrderedDict
# 从common.optimizer导入所有优化器
from common.optimizer import *

# 定义优化目标函数
def f(x, y):
    """目标函数：f(x,y) = x²/20 + y²"""
    return x**2 / 20.0 + y**2

# 定义目标函数的梯度
def df(x, y):
    """计算目标函数的梯度"""
    return x / 10.0, 2.0*y  # df/dx = x/10, df/dy = 2y

# 设置初始参数位置
init_pos = (-7.0, 2.0)  # 初始点(x,y)
params = {}  # 参数字典
params['x'], params['y'] = init_pos[0], init_pos[1]  # 初始化参数
grads = {}  # 梯度字典
grads['x'], grads['y'] = 0, 0  # 初始化梯度

# 初始化各种优化器(使用有序字典保持顺序)
optimizers = OrderedDict()
optimizers["SGD"] = SGD(lr=0.95)          # 随机梯度下降
optimizers["Momentum"] = Momentum(lr=0.1) # 动量法
optimizers["AdaGrad"] = AdaGrad(lr=1.5)   # AdaGrad
optimizers["Adam"] = Adam(lr=0.3)         # Adam

# 绘图索引
idx = 1

# 对每个优化器进行测试
for key in optimizers:
    optimizer = optimizers[key]  # 获取当前优化器
    x_history = []  # 记录x的历史值
    y_history = []  # 记录y的历史值
    # 重置参数到初始位置
    params['x'], params['y'] = init_pos[0], init_pos[1]

    # 优化迭代(30次)
    for i in range(30):
        # 记录当前参数位置
        x_history.append(params['x'])
        y_history.append(params['y'])

        # 计算梯度
        grads['x'], grads['y'] = df(params['x'], params['y'])
        # 更新参数
        optimizer.update(params, grads)

    # 准备绘图数据
    x = np.arange(-10, 10, 0.01)  # x轴范围
    y = np.arange(-5, 5, 0.01)    # y轴范围

    # 生成网格点
    X, Y = np.meshgrid(x, y)
    # 计算各网格点的函数值
    Z = f(X, Y)

    # 简化等高线显示(去掉高值区域)
    mask = Z > 7
    Z[mask] = 0

    # 绘制子图(2行2列)
    plt.subplot(2, 2, idx)
    idx += 1
    # 绘制优化路径(红色圆点连线)
    plt.plot(x_history, y_history, 'o-', color="red")
    # 绘制等高线
    plt.contour(X, Y, Z)
    # 设置坐标轴范围
    plt.ylim(-10, 10)
    plt.xlim(-10, 10)
    # 标记最小值点(0,0)
    plt.plot(0, 0, '+')
    # 设置标题和坐标轴标签
    plt.title(key)
    plt.xlabel("x")
    plt.ylabel("y")

# 显示图形
plt.show()

"""
预期结果分析：
1. SGD: 沿梯度方向直线下降，在y方向震荡明显
2. Momentum: 在y方向震荡减弱，有"过冲"现象
3. AdaGrad: 自适应调整学习率，在y方向快速收敛
4. Adam: 结合动量和自适应学习率，路径最平滑

注意事项：
1. 不同优化器的学习率需要单独调整
2. 迭代次数影响最终收敛位置
3. 函数形状影响优化难度(y方向更陡峭)
4. 实际应用中需要根据问题选择合适的优化器
"""
